import os
import os.path as osp
from Models.models.irpe import get_rpe_config
from Models.utils import ensure_path

def print_save_path(args):
    # pretrain阶段使用
    # 打印储存空间
    # 使用可变变量引用传参，不用显示赋值
    args.train_type = "epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}".format(
        lr=args.lr, stepsize=args.step_size, gamma=args.gamma, imagesize=args.image_size, optim=args.optim, epoch=args.max_epoch)
    if args.model == "resnet":
        if args.with_SA:
            if args.no_mlp and not args.SA_res:
                args.model_type = "{model}_MySA({heads}_{dim_head})".format(
                    model=args.model, heads=args.SA_heads, dim_head=args.SA_dim_head)
            elif args.no_mlp and args.SA_res:
                args.model_type = "{model}_MyResSA({heads}_{dim_head})".format(
                    model=args.model, heads=args.SA_heads, dim_head=args.SA_dim_head)
            else:          
                args.model_type = "{model}_SA({depth}_{heads}_{dim_head}_{mlp_dim})".format(
                    model=args.model, heads=args.SA_heads, dim_head=args.SA_dim_head, mlp_dim=args.SA_mlp_dim, depth=args.SA_depth)
            if args.pos_embed:
                args.model_type += "_pos-embed"
            elif args.use_rpe:
                args.model_type += "_use-rpe" + "_{}_{}_{}_{}".format(args.rpe_method, args.rpe_mode, args.rpe_on, args.rpe_ratio)
            elif args.use_origin_outlook:
                if args.SA_res:
                    args.model_type = "{model}_MyResSA_outlook({heads}_{kernel})".format(
                        model=args.model, heads=args.SA_heads, kernel=args.outlook_kernel)
                else:
                    args.model_type = "{model}_MySA_outlook({heads}_{kernel})".format(
                        model=args.model, heads=args.SA_heads, kernel=args.outlook_kernel)      
            elif args.use_involution:
                if args.SA_res:
                    args.model_type = "{model}_MyResSA_involution({kernel}_softmax-{invo_softmax})".format(
                        model=args.model, kernel=args.involution_kernel, invo_softmax=args.involution_softmax)
                else:
                    args.model_type = "{model}_MySA_involution({kernel}_softmax-{invo_softmax})".format(
                        model=args.model, kernel=args.involution_kernel, invo_softmax=args.involution_softmax)
            elif args.use_carafe_attn:
                if args.SA_res:
                    args.model_type = "{model}_MyResSA_carafe({kernel})".format(
                            model=args.model, kernel=args.carafe_kernel)
                else:
                    args.model_type = "{model}_MySA_carafe({kernel})".format(
                            model=args.model, kernel=args.carafe_kernel)
            if args.SA_relu:
                args.model_type += "_relu"
            args.save_path = "pre_train/{dataset}/".format(dataset=args.dataset) + args.model_type + "_" + args.train_type

        else:
            args.save_path = 'pre_train/{dataset}/{model}_epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}'.format(
                dataset=args.dataset, model=args.model, lr=args.lr, stepsize=args.step_size, gamma=args.gamma, imagesize=args.image_size, optim=args.optim, epoch=args.max_epoch)
        
        if args.use_deformconv and args.modulated_deformconv:
            args.save_path += "_modulated_deform"
        elif args.use_deformconv and not args.modulated_deformconv:
            args.save_path += "_deform"
        if args.MLA:
            layers = ""
            for i in args.MLA:
                layers += str(i)
            args.save_path += "_MLA{}".format(layers)
        if args.use_elu_taylor:
            args.save_path += '_ET'
        if args.use_taylor:
            args.save_path += '_Taylor'
        if args.use_sig_softmax:
            args.save_path += '_SigSoftmax'
        if args.use_focal_loss:
            args.save_path += '_FocalLoss'
        if args.metric == "l2":
            args.save_path += "_l2"

        args.save_path += "_" + args.sche

    elif args.model == "ViT":
        args.save_path = 'pre_train/{dataset}/{model}_depth{depth}_epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}_use-clstoken({class_token})_vit-mode({vit_mode})'.format(
            dataset=args.dataset, model=args.model, lr=args.lr, stepsize=args.step_size, gamma=args.gamma, imagesize=args.image_size, class_token=str(not args.not_use_clstoken), vit_mode=args.vit_mode, optim=args.optim,
            epoch=args.max_epoch, depth=args.vit_depth)

    elif args.model == "vit_small_patch16_224":
        args.save_path = 'pre_train/{dataset}/{model}_depth{depth}_epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}_use-imagenet-params({imagenet_pretrain}))'.format(
            dataset=args.dataset, model=args.model, lr=args.lr, stepsize=args.step_size, gamma=args.gamma, imagesize=args.image_size, optim=args.optim, epoch=args.max_epoch, depth=args.vit_depth, imagenet_pretrain=str(not args.not_imagenet_pretrain))

    args.save_path = osp.join('checkpoint', args.save_path)
    if args.extra_dir is not None:
        args.save_path = osp.join(args.save_path, args.extra_dir)
    ensure_path(args.save_path)
    return args.save_path


def pretrain_save_path(args):
    # train_meta阶段使用，找寻pretrain model的存储地址
    args.train_type = "epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}".format(
        lr=args.pre_lr, stepsize=args.pre_step_size, gamma=args.pre_gamma, imagesize=args.image_size, optim=args.pre_optim, epoch=args.pre_epoch)
    if args.model == "resnet":
        if args.with_SA:
            if args.no_mlp and not args.SA_res:
                args.model_type = "{model}_MySA({heads}_{dim_head})".format(
                    model=args.model, heads=args.SA_heads, dim_head=args.SA_dim_head)
            elif args.no_mlp and args.SA_res:
                args.model_type = "{model}_MyResSA({heads}_{dim_head})".format(
                    model=args.model, heads=args.SA_heads, dim_head=args.SA_dim_head)
            else:          
                args.model_type = "{model}_SA({depth}_{heads}_{dim_head}_{mlp_dim})".format(
                    model=args.model, heads=args.SA_heads, dim_head=args.SA_dim_head, mlp_dim=args.SA_mlp_dim, depth=args.SA_depth)
            if args.pos_embed:
                args.model_type += "_pos-embed"
            elif args.use_rpe:
                args.model_type += "_use-rpe" + "_{}_{}_{}_{}".format(args.rpe_method, args.rpe_mode, args.rpe_on, args.rpe_ratio)
            elif args.use_origin_outlook:
                if args.SA_res:
                    args.model_type = "{model}_MyResSA_outlook({heads}_{kernel})".format(
                        model=args.model, heads=args.SA_heads, kernel=args.outlook_kernel)
                else:
                    args.model_type = "{model}_MySA_outlook({heads}_{kernel})".format(
                        model=args.model, heads=args.SA_heads, kernel=args.outlook_kernel)      
            elif args.use_involution:
                if args.SA_res:
                    args.model_type = "{model}_MyResSA_involution({kernel}_softmax-{invo_softmax})".format(
                        model=args.model, kernel=args.involution_kernel, invo_softmax=args.involution_softmax)
                else:
                    args.model_type = "{model}_MySA_involution({kernel}_softmax-{invo_softmax})".format(
                        model=args.model, kernel=args.involution_kernel, invo_softmax=args.involution_softmax)
            elif args.use_carafe_attn:
                if args.SA_res:
                    args.model_type = "{model}_MyResSA_carafe({kernel})".format(
                            model=args.model, kernel=args.carafe_kernel)
                else:
                    args.model_type = "{model}_MySA_carafe({kernel})".format(
                            model=args.model, kernel=args.carafe_kernel)
            if args.SA_relu:
                args.model_type += "_relu"
            args.pre_save_path = "pre_train/{dataset}/".format(dataset=args.dataset) + args.model_type + "_" + args.train_type
        else:
            args.pre_save_path = 'pre_train/{dataset}/{model}_epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}'.format(
                dataset=args.dataset, model=args.model, lr=args.pre_lr, stepsize=args.pre_step_size, gamma=args.pre_gamma, imagesize=args.image_size, optim=args.pre_optim, epoch=args.pre_epoch)
        
        if args.use_deformconv and args.modulated_deformconv:
            args.pre_save_path += "_modulated_deform"
        elif args.use_deformconv and not args.modulated_deformconv:
            args.pre_save_path += "_deform"
        if args.MLA:
            layers = ""
            for i in args.MLA:
                layers += str(i)
            args.pre_save_path += "_MLA{}".format(layers)
        if args.use_elu_taylor:
            args.pre_save_path += '_ET'
        if args.use_taylor:
            args.pre_save_path += '_Taylor'
        if args.use_sig_softmax:
            args.pre_save_path += '_SigSoftmax'
        if args.use_focal_loss:
            args.pre_save_path += '_FocalLoss'
        if args.metric == "l2":
            args.pre_save_path += "_l2"
        
        args.pre_save_path += "_" + args.sche

    elif args.model == "ViT":
        args.pre_save_path = 'pre_train/{dataset}/{model}_depth{depth}_epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}_use-clstoken({class_token})_vit-mode({vit_mode})'.format(
            dataset=args.dataset, model=args.model, lr=args.pre_lr, stepsize=args.pre_step_size, gamma=args.pre_gamma, imagesize=args.image_size, class_token=str(not args.not_use_clstoken), vit_mode=args.vit_mode, optim=args.pre_optim,
            epoch=args.pre_epoch, depth=args.vit_depth)

    elif args.model == "vit_small_patch16_224":
        args.pre_save_path = 'pre_train/{dataset}/{model}_depth{depth}_epoch{epoch}_optim{optim}_lr{lr:.4f}_stepsize{stepsize}_gamma{gamma:.2f}_imagesize{imagesize}_use-imagenet-params({imagenet_pretrain}))'.format(
            dataset=args.dataset, model=args.model, lr=args.pre_lr, stepsize=args.pre_step_size, gamma=args.pre_gamma, imagesize=args.image_size, optim=args.pre_optim, epoch=args.pre_epoch, depth=args.vit_depth, imagenet_pretrain=str(not args.not_imagenet_pretrain))

    args.pre_save_path = osp.join('checkpoint', args.pre_save_path)
    if args.extra_dir is not None:
        args.pre_save_path = osp.join(args.pre_save_path, args.extra_dir)
    if os.path.exists(args.pre_save_path):
        print("预训练模型路径:{}".format(args.pre_save_path))
    else:
        raise ValueError("没有该路径:{}".format(args.pre_save_path))
    return args.pre_save_path

def parse_tune_pretrain(args):
    if args.model == "ViT" and args.deepemd != 'fcn':
        print("选用ViT时未将deepemd参数置为fcn模式,当前为{}模式,将转换为fcn模式".format(args.deepemd))
        args.deepemd = 'fcn'

    if args.model == "ViT" and args.image_size != 256:
        print("选用ViT时未将image_size调整为256, 当前image size = {},将转换为256".format(
            args.image_size))
        args.image_size = 256

    if args.model == "vit_small_patch16_224" and args.image_size != 224:
        print("选用vit_small_patch16_224时未将image_size调整为224, 当前image size = {},将转换为224".format(
            args.image_size))
        args.image_size = 224

    if args.model == 'resnet' and args.image_size != 84:
        print("选用resnet时未将image_size调整为84, 当前image size = {},将转换为84".format(
            args.image_size))
        args.image_size = 84

    if args.model == 'resnet' and args.with_SA:
        print("使用带self attention的resnet")
        
    if args.MLA:
        args.MLA = [int(x) for x in args.MLA.split(',')]
        assert len(args.MLA) <= 4, "MLA最大层数为4, 因为BACKBONE只有四层输出"
        for i in args.MLA:
            assert i <= 3, "MLA制定层数不可超过3"

    if args.use_rpe:
        rpe_config = get_rpe_config(ratio=args.rpe_ratio, method=args.rpe_method, mode=args.rpe_mode, rpe_on=args.rpe_on)
        args.rpe_config = rpe_config
        


def meta_save_path(args):
    # meta train阶段使用
    epoch_index = args.pre_save_path.find("_epoch")
    args.model_name = args.pre_save_path[:epoch_index].split("/")[-1]
    if args.use_deformconv and args.modulated_deformconv:
        args.model_name += "_modulated_deform"
    elif args.use_deformconv and not args.modulated_deformconv:
        args.model_name += "_deform"
    if args.MLA:
        layers = ""
        for i in args.MLA:
            layers += str(i)
        args.model_name += "_MLA{}".format(layers)
    if args.use_elu_taylor:
        args.model_name += '_ET'
    if args.use_taylor:
        args.model_name += '_Taylor'
    if args.use_sig_softmax:
        args.model_name += '_SigSoftmax'
    if args.use_focal_loss:
        args.model_name += '_FocalLoss'
    if args.metric == "l2":
        args.model_name += "_l2"

    if args.sfc_update_step == 100:
        args.save_path = "{dataset}/{model_name}/{shot}shot-{way}way".format(dataset=args.dataset,
                        model_name=args.model_name, shot=args.shot, way=args.way, sfc_update_step=int(args.sfc_update_step))
    else:
        args.save_path = "{dataset}/{model_name}/{shot}shot-{way}way_SFC{sfc_update_step}".format(dataset=args.dataset,
                        model_name=args.model_name, shot=args.shot, way=args.way, sfc_update_step=int(args.sfc_update_step))
    args.save_path = osp.join('checkpoint/meta_train',
                              args.save_path + "_{}".format(args.solver))
    if args.extra_dir is not None:
        args.save_path = osp.join(args.save_path, args.extra_dir)
        
    args.save_path += "_" + args.sche
    ensure_path(args.save_path)
    return args.save_path, args.model_name


def format_model_name(args):
    if args.model_name == "resnet_MyResSA" or "resnet_MySA":
        args.model_name = args.model_name + \
            "({}_{})".format(args.SA_heads, args.SA_dim_head)
        if args.pos_embed:
            args.model_name += "_pos-embed"
        elif args.use_rpe:
            args.model_name += "_use-rpe" + "_{}_{}_{}_{}".format(args.rpe_method, args.rpe_mode, args.rpe_on, args.rpe_ratio)
        elif args.use_origin_outlook:
            if args.SA_res:
                args.model_name = "{model}_MyResSA_outlook({heads}_{kernel})".format(
                    model=args.model, heads=args.SA_heads, kernel=args.outlook_kernel)
            else:
                args.model_name = "{model}_MySA_outlook({heads}_{kernel})".format(
                    model=args.model, heads=args.SA_heads, kernel=args.outlook_kernel)

        elif args.use_involution:
            if args.SA_res:
                args.model_name = "{model}_MyResSA_involution({kernel}_softmax-{invo_softmax})".format(
                        model=args.model, kernel=args.involution_kernel, invo_softmax=args.involution_softmax)
            else:
                args.model_name = "{model}_MySA_involution({kernel}_softmax-{invo_softmax})".format(
                        model=args.model, kernel=args.involution_kernel, invo_softmax=args.involution_softmax)
        
        elif args.use_carafe_attn:
            if args.SA_res:
                args.model_name = "{model}_MyResSA_carafe({kernel})".format(
                        model=args.model, kernel=args.carafe_kernel)
            else:
                args.model_name = "{model}_MySA_carafe({kernel})".format(
                        model=args.model, kernel=args.carafe_kernel)
        if args.SA_relu:
            args.model_name += "_relu"

        if args.MLA:
            layers = ""
            for i in args.MLA:
                layers += str(i)
            args.model_name += "_MLA{}".format(layers)

        if args.use_elu_taylor:
            args.model_name += '_ET'
        if args.use_taylor:
            args.model_name += '_Taylor'
        if args.use_sig_softmax:
            args.model_name += '_SigSoftmax'
        if args.use_focal_loss:
            args.model_name += '_FocalLoss'
        if args.metric == "l2":
            args.model_name += "_l2"
        
    elif args.model_name == "resnet":
        pass

    else:
        # TODO: 实现其他方法
        print("")
        raise ValueError("没有该model_name")

def eval_res_save_path(args):
    if args.origin:
        args.res_save_path = "result/{dataset}/resnet_origin/{shot}shot-{way}way/".format(
            dataset=args.dataset, shot=args.shot, way=args.way)

    else :
        if args.sfc_update_step == 100:
            if args.use_specific_status:
                args.res_save_path = "result/{dataset}/{model_name}/{shot}shot-{way}way_{sche}_use_specific_status/".format(
                    dataset=args.dataset, model_name=args.model_name, shot=args.shot, way=args.way, sche=args.sche)
            else:
                args.res_save_path = "result/{dataset}/{model_name}/{shot}shot-{way}way_{sche}/".format(
                    dataset=args.dataset, model_name=args.model_name, shot=args.shot, way=args.way, sche=args.sche)
        else:
            if args.use_specific_status:
                args.res_save_path = "result/{dataset}/{model_name}/{shot}shot-{way}way_{sche}_SFC{sfc_update_step}_use_specific_status/".format(
                    dataset=args.dataset, model_name=args.model_name, shot=args.shot, way=args.way, sfc_update_step=args.sfc_update_step, sche=args.sche)
            else:
                args.res_save_path = "result/{dataset}/{model_name}/{shot}shot-{way}way_{sche}_SFC{sfc_update_step}/".format(
                    dataset=args.dataset, model_name=args.model_name, shot=args.shot, way=args.way, sfc_update_step=args.sfc_update_step, sche=args.sche)
